"""
Basis functions module for DLR-FEM implementation.

This module provides polynomial basis functions for finite element methods,
specifically Lagrange polynomials on Gauss-Lobatto points.
"""

import numpy as np
from typing import Tuple, List, Callable
from .mesh import gauss_lobatto_points


class LagrangeBasis:
    """
    Lagrange polynomial basis functions on Gauss-Lobatto points.
    
    This class provides evaluation of Lagrange basis functions and their
    derivatives, which are used in spectral element methods.
    """
    
    def __init__(self, poly_order: int):
        """
        Initialize Lagrange basis.
        
        Args:
            poly_order: Polynomial order (degree of basis functions)
        """
        self.poly_order = poly_order
        self.n_points = poly_order + 1
        
        # Get Gauss-Lobatto points (nodes)
        self.nodes, self.weights = gauss_lobatto_points(self.n_points)
        
        # Precompute derivative matrix
        self.D = self._compute_derivative_matrix()
    
    def evaluate(self, xi: np.ndarray, i: int) -> np.ndarray:
        """
        Evaluate i-th Lagrange basis function at points xi.
        
        Args:
            xi: Evaluation points in reference domain [-1, 1]
            i: Basis function index (0 <= i < n_points)
            
        Returns:
            Values of i-th basis function at xi
        """
        if i < 0 or i >= self.n_points:
            raise ValueError(f"Basis function index {i} out of range [0, {self.n_points-1}]")
        
        xi = np.asarray(xi)
        result = np.ones_like(xi, dtype=float)
        
        for j in range(self.n_points):
            if j != i:
                result *= (xi - self.nodes[j]) / (self.nodes[i] - self.nodes[j])
        
        return result
    
    def evaluate_derivative(self, xi: np.ndarray, i: int) -> np.ndarray:
        """
        Evaluate derivative of i-th Lagrange basis function at points xi.
        
        Args:
            xi: Evaluation points in reference domain [-1, 1]
            i: Basis function index
            
        Returns:
            Derivative values of i-th basis function at xi
        """
        if i < 0 or i >= self.n_points:
            raise ValueError(f"Basis function index {i} out of range [0, {self.n_points-1}]")
        
        xi = np.asarray(xi)
        result = np.zeros_like(xi, dtype=float)
        
        # Sum over all terms in the product rule
        for k in range(self.n_points):
            if k != i:
                term = np.ones_like(xi, dtype=float)
                
                # Product of all factors except the k-th
                for j in range(self.n_points):
                    if j != i and j != k:
                        term *= (xi - self.nodes[j]) / (self.nodes[i] - self.nodes[j])
                
                # Derivative contribution from k-th factor
                term /= (self.nodes[i] - self.nodes[k])
                result += term
        
        return result
    
    def evaluate_all(self, xi: np.ndarray) -> np.ndarray:
        """
        Evaluate all basis functions at points xi.
        
        Args:
            xi: Evaluation points in reference domain [-1, 1]
            
        Returns:
            Array of shape (len(xi), n_points) with basis function values
        """
        xi = np.asarray(xi)
        if xi.ndim == 0:
            xi = xi.reshape(1)
        
        result = np.zeros((len(xi), self.n_points))
        
        for i in range(self.n_points):
            result[:, i] = self.evaluate(xi, i)
        
        return result
    
    def evaluate_all_derivatives(self, xi: np.ndarray) -> np.ndarray:
        """
        Evaluate derivatives of all basis functions at points xi.
        
        Args:
            xi: Evaluation points in reference domain [-1, 1]
            
        Returns:
            Array of shape (len(xi), n_points) with derivative values
        """
        xi = np.asarray(xi)
        if xi.ndim == 0:
            xi = xi.reshape(1)
        
        result = np.zeros((len(xi), self.n_points))
        
        for i in range(self.n_points):
            result[:, i] = self.evaluate_derivative(xi, i)
        
        return result
    
    def _compute_derivative_matrix(self) -> np.ndarray:
        """
        Compute derivative matrix D_ij = d/dxi L_j(xi_i).
        
        This matrix allows efficient computation of derivatives at the nodes.
        """
        D = np.zeros((self.n_points, self.n_points))
        
        for i in range(self.n_points):
            for j in range(self.n_points):
                if i != j:
                    # Off-diagonal terms
                    numerator = 1.0
                    denominator = 1.0
                    
                    for k in range(self.n_points):
                        if k != j:
                            numerator *= (self.nodes[i] - self.nodes[k])
                        if k != i:
                            denominator *= (self.nodes[j] - self.nodes[k])
                    
                    D[i, j] = numerator / denominator
                else:
                    # Diagonal terms: sum of 1/(xi_i - xi_k) for k != i
                    D[i, i] = 0.0
                    for k in range(self.n_points):
                        if k != i:
                            D[i, i] += 1.0 / (self.nodes[i] - self.nodes[k])
        
        return D
    
    def interpolate_to_points(self, values_at_nodes: np.ndarray, 
                            xi: np.ndarray) -> np.ndarray:
        """
        Interpolate function values from nodes to arbitrary points.
        
        Args:
            values_at_nodes: Function values at Gauss-Lobatto nodes
            xi: Target interpolation points
            
        Returns:
            Interpolated values at xi
        """
        if len(values_at_nodes) != self.n_points:
            raise ValueError(f"Expected {self.n_points} values, got {len(values_at_nodes)}")
        
        basis_values = self.evaluate_all(xi)
        return basis_values @ values_at_nodes
    
    def compute_nodal_derivatives(self, values_at_nodes: np.ndarray) -> np.ndarray:
        """
        Compute derivative values at nodes using derivative matrix.
        
        Args:
            values_at_nodes: Function values at Gauss-Lobatto nodes
            
        Returns:
            Derivative values at nodes
        """
        if len(values_at_nodes) != self.n_points:
            raise ValueError(f"Expected {self.n_points} values, got {len(values_at_nodes)}")
        
        return self.D @ values_at_nodes


class TensorProductBasis:
    """
    2D tensor product basis using 1D Lagrange basis functions.
    
    This provides the basis functions Φ(x) ⊗ Ψ(y) for the DLR-FEM method.
    """
    
    def __init__(self, poly_order_x: int, poly_order_y: int):
        """
        Initialize tensor product basis.
        
        Args:
            poly_order_x: Polynomial order in x-direction
            poly_order_y: Polynomial order in y-direction
        """
        self.poly_order_x = poly_order_x
        self.poly_order_y = poly_order_y
        
        self.basis_x = LagrangeBasis(poly_order_x)
        self.basis_y = LagrangeBasis(poly_order_y)
        
        self.nx = self.basis_x.n_points
        self.ny = self.basis_y.n_points
        
        # Total number of 2D basis functions
        self.n_basis = self.nx * self.ny
    
    def evaluate_2d(self, x: np.ndarray, y: np.ndarray, 
                   i: int, j: int) -> np.ndarray:
        """
        Evaluate 2D basis function Φ_i(x) * Ψ_j(y).
        
        Args:
            x, y: Evaluation points
            i: Basis function index in x-direction (0 <= i < nx)
            j: Basis function index in y-direction (0 <= j < ny)
            
        Returns:
            Values of 2D basis function
        """
        phi_x = self.basis_x.evaluate(x, i)
        psi_y = self.basis_y.evaluate(y, j)
        return phi_x * psi_y
    
    def evaluate_gradient_2d(self, x: np.ndarray, y: np.ndarray,
                           i: int, j: int) -> Tuple[np.ndarray, np.ndarray]:
        """
        Evaluate gradient of 2D basis function.
        
        Args:
            x, y: Evaluation points
            i, j: Basis function indices
            
        Returns:
            (grad_x, grad_y): Gradient components
        """
        phi_x = self.basis_x.evaluate(x, i)
        phi_x_deriv = self.basis_x.evaluate_derivative(x, i)
        psi_y = self.basis_y.evaluate(y, j)
        psi_y_deriv = self.basis_y.evaluate_derivative(y, j)
        
        grad_x = phi_x_deriv * psi_y
        grad_y = phi_x * psi_y_deriv
        
        return grad_x, grad_y
    
    def get_phi_matrix(self, x_points: np.ndarray) -> np.ndarray:
        """
        Get matrix Φ(x) where Φ_ij = φ_j(x_i).
        
        Args:
            x_points: Evaluation points in x-direction
            
        Returns:
            Matrix of shape (len(x_points), nx)
        """
        return self.basis_x.evaluate_all(x_points)
    
    def get_psi_matrix(self, y_points: np.ndarray) -> np.ndarray:
        """
        Get matrix Ψ(y) where Ψ_ij = ψ_j(y_i).
        
        Args:
            y_points: Evaluation points in y-direction
            
        Returns:
            Matrix of shape (len(y_points), ny)
        """
        return self.basis_y.evaluate_all(y_points)
    
    def get_phi_derivative_matrix(self, x_points: np.ndarray) -> np.ndarray:
        """
        Get derivative matrix ∂Φ/∂x.
        
        Args:
            x_points: Evaluation points in x-direction
            
        Returns:
            Derivative matrix of shape (len(x_points), nx)
        """
        return self.basis_x.evaluate_all_derivatives(x_points)
    
    def get_psi_derivative_matrix(self, y_points: np.ndarray) -> np.ndarray:
        """
        Get derivative matrix ∂Ψ/∂y.
        
        Args:
            y_points: Evaluation points in y-direction
            
        Returns:
            Derivative matrix of shape (len(y_points), ny)
        """
        return self.basis_y.evaluate_all_derivatives(y_points)


def create_basis_functions(poly_order: int) -> TensorProductBasis:
    """
    Convenience function to create tensor product basis with same order in both directions.
    
    Args:
        poly_order: Polynomial order for both x and y directions
        
    Returns:
        TensorProductBasis instance
    """
    return TensorProductBasis(poly_order, poly_order)


if __name__ == "__main__":
    # Test Lagrange basis
    print("Testing Lagrange basis functions...")
    basis = LagrangeBasis(poly_order=2)
    
    print(f"Nodes: {basis.nodes}")
    print(f"Weights: {basis.weights}")
    
    # Test partition of unity
    xi_test = np.linspace(-1, 1, 10)
    all_basis = basis.evaluate_all(xi_test)
    unity_check = np.sum(all_basis, axis=1)
    print(f"Partition of unity check (should be all 1s): {unity_check[:3]}...")
    
    # Test tensor product basis
    print("\nTesting tensor product basis...")
    tensor_basis = create_basis_functions(poly_order=2)
    print(f"2D basis has {tensor_basis.n_basis} functions")
    print(f"Shape: {tensor_basis.nx} × {tensor_basis.ny}")
    
    # Test evaluation
    x_test = np.array([0.0])
    y_test = np.array([0.0])
    val = tensor_basis.evaluate_2d(x_test, y_test, 1, 1)  # Central basis function
    print(f"Central basis function at origin: {val[0]:.4f}")